use anyhow::Result;
use ipnetwork::IpNetwork;
use std::net::IpAddr;
use std::sync::mpsc::{channel, Sender, Receiver};
use std::str::FromStr;
use surge_ping::{Client, Config, PingIdentifier, PingSequence, ICMP};
use tokio;
use once_cell::sync::Lazy;
use log::{debug, error, info};
use futures::stream::{self, StreamExt};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;

// 命令枚举
#[derive(Debug)]
pub enum Command {
    Help,
    Scan,
    Reset,
    Quit,
    SetNetwork(IpNetwork),
}

#[derive(Debug, Clone)]
pub struct IpStatus {
    pub address: String,
    pub status: Status,
}

#[derive(Debug, Clone)]
pub enum Status {
    Unknown,
    Pending,
    Success(u64),  // 延迟时间（毫秒）
    Timeout,
}

#[derive(Debug)]
enum UpdateMessage {
    StatusUpdate(String, Status),  // IP地址和状态
    Progress(usize, usize),  // 添加进度信息 (当前, 总数)
    ScanComplete,                  // 扫描完成
    Error(String),                 // 错误信息
}

pub struct App {
    pub input: String,
    pub network: IpNetwork,
    pub ip_addresses: Vec<IpStatus>,
    pub message: String,
    pub message_scroll: u16,
    pub scanning: bool,
    current_progress: usize,
    total_progress: usize,
    update_rx: Receiver<UpdateMessage>,
    update_tx: Sender<UpdateMessage>,
    start_time: std::time::Instant,
}

const HELP_TEXT: &str = "\
可用命令:
  h, help           显示帮助信息
  s, scan           开始扫描
  r, reset          重置扫描
  q, quit           退出程序
  <network/mask>    设置网段 (例如: 192.168.1.0/24)

说明:
- 输入命令后按回车执行
- 扫描时会显示每个IP的延迟时间
- 无响应的IP会标记为 timeout";

const COLOR_LEGEND: &str = "\
颜色说明:
  绿色     - 延迟 < 10ms
  浅绿色   - 延迟 < 50ms
  黄绿色   - 延迟 < 100ms
  黄色     - 延迟 < 500ms
  橙色     - 延迟 < 600ms
  红色     - 超时
  灰色     - 未扫描";

impl App {
    pub fn new(network: IpNetwork) -> Self {
        let (tx, rx) = channel();
        let mut app = Self {
            input: String::new(),
            network,
            ip_addresses: Vec::new(),
            message: String::from("Type 'h' or 'help' for available commands"),
            message_scroll: 0,
            scanning: false,
            current_progress: 0,
            total_progress: 0,
            update_rx: rx,
            update_tx: tx,
            start_time: std::time::Instant::now(),
        };
        app.ip_addresses = app.generate_ip_list(&network);
        app
    }

    fn generate_ip_list(&self, network: &IpNetwork) -> Vec<IpStatus> {
        debug!("Generating IP list for network: {}", network);
        let mut ip_addresses = Vec::new();
        
        // 确保网络前缀不超过31
        if network.prefix() >= 32 {
            error!("Invalid network prefix: /{}", network.prefix());
            return ip_addresses;
        }

        let size = u32::from(network.prefix()) as usize;
        let num_hosts = if size < 32 { 1 << (32 - size) } else { 0 };
        
        debug!("Network size: {}, Number of hosts: {}", size, num_hosts);
        
        // 跳过网络地址和广播地址
        for ip in network.iter().skip(1).take(num_hosts - 2) {
            ip_addresses.push(IpStatus {
                address: ip.to_string(),
                status: Status::Unknown,
            });
        }
        
        debug!("Generated {} IP addresses", ip_addresses.len());
        ip_addresses
    }

    pub fn start_scan(&mut self) -> Result<()> {
        if self.scanning {
            info!("Scan already in progress");
            self.message = "Scan already in progress".to_string();
            return Ok(());
        }

        info!("Starting network scan for {}", self.network);
        self.scanning = true;
        self.message = "Scanning network...".to_string();
        self.current_progress = 0;
        self.total_progress = self.ip_addresses.len();

        // 重置所有 IP 状态
        for ip in &mut self.ip_addresses {
            ip.status = Status::Pending;
        }

        let tx = self.update_tx.clone();
        let addresses = self.ip_addresses.clone();
        let total = self.total_progress;
        let network = self.network.to_string();
        let current_progress = Arc::new(AtomicUsize::new(0));

        static RT: Lazy<tokio::runtime::Runtime> = Lazy::new(|| {
            tokio::runtime::Runtime::new().unwrap()
        });

        let start_time = std::time::Instant::now();  // 添加计时器
        info!("Starting scan with {} IPs, timeout: 600ms, concurrency: 64", self.ip_addresses.len());

        RT.spawn(async move {
            debug!("Creating ICMP client for network {}", network);
            let config = Config::builder()
                .kind(ICMP::V4)
                .build();
            let client = match Client::new(&config) {
                Ok(client) => client,
                Err(e) => {
                    error!("Failed to create ICMP client: {}", e);
                    let _ = tx.send(UpdateMessage::Error(e.to_string()));
                    return Ok::<_, anyhow::Error>(());
                }
            };

            stream::iter(addresses)
                .map(|ip| {
                    let client = &client;
                    let tx = tx.clone();
                    let progress = current_progress.clone();
                    
                    async move {
                        if let Ok(addr) = ip.address.parse::<IpAddr>() {
                            let mut pinger = client.pinger(addr, PingIdentifier(1)).await;
                            pinger.timeout(Duration::from_millis(600));  // 修改为600ms
                            
                            match pinger.ping(PingSequence(1), &[0; 16]).await {
                                Ok((_, duration)) => {
                                    let ms = duration.as_millis() as u64;
                                    debug!("Response from {} in {}ms", addr, ms);
                                    let _ = tx.send(UpdateMessage::StatusUpdate(
                                        ip.address,
                                        Status::Success(ms)
                                    ));
                                }
                                Err(e) => {
                                    debug!("Timeout for {}", addr);
                                    let _ = tx.send(UpdateMessage::StatusUpdate(
                                        ip.address,
                                        Status::Timeout
                                    ));
                                }
                            }
                        }
                        
                        let current = progress.fetch_add(1, Ordering::SeqCst) + 1;
                        let _ = tx.send(UpdateMessage::Progress(current, total));
                    }
                })
                .buffer_unordered(64)
                .collect::<Vec<_>>()
                .await;

            let elapsed = start_time.elapsed();
            info!("Scan completed for network {} in {:.2?}", network, elapsed);
            let _ = tx.send(UpdateMessage::ScanComplete);
            Ok::<_, anyhow::Error>(())
        });

        Ok(())
    }

    pub fn progress_ratio(&self) -> f64 {
        if self.total_progress == 0 {
            0.0
        } else {
            self.current_progress as f64 / self.total_progress as f64
        }
    }

    pub fn update(&mut self) {
        while let Ok(msg) = self.update_rx.try_recv() {
            match msg {
                UpdateMessage::StatusUpdate(addr, status) => {
                    if let Some(ip) = self.ip_addresses.iter_mut()
                        .find(|ip| ip.address == addr) {
                        ip.status = status;
                    }
                }
                UpdateMessage::Progress(current, total) => {
                    self.current_progress = current;
                    self.total_progress = total;
                    self.message = format!("Scanning... {}/{} IPs", current, total);
                }
                UpdateMessage::ScanComplete => {
                    self.scanning = false;
                    let elapsed = self.start_time.elapsed();
                    self.message = format!("扫描完成，耗时: {:.2?}\n\n{}", elapsed, COLOR_LEGEND);
                }
                UpdateMessage::Error(err) => {
                    self.scanning = false;
                    self.message = format!("Error: {}", err);
                }
            }
        }
    }

    pub fn handle_command(&mut self) -> Result<Option<Command>> {
        let command = self.parse_command();
        self.message_scroll = 0;
        match command {
            Some(Command::Help) => {
                self.message = HELP_TEXT.to_string();
            }
            Some(Command::SetNetwork(new_network)) => {
                self.network = new_network;
                self.ip_addresses = self.generate_ip_list(&new_network);
                self.message = format!("Network set to: {}", new_network);
            }
            Some(Command::Scan) => {
                self.start_scan()?;
            }
            Some(Command::Reset) => {
                self.ip_addresses = self.generate_ip_list(&self.network);
                self.message = "Reset scan results".to_string();
            }
            None => {
                self.message = format!("Unknown command: '{}'\n{}", self.input, HELP_TEXT);
            }
            _ => {}
        }
        self.input.clear();
        Ok(command)
    }

    pub fn parse_command(&self) -> Option<Command> {
        let input = self.input.trim();
        debug!("Parsing command input: {}", input);
        
        match input {
            "h" | "help" => Some(Command::Help),
            "q" | "quit" => Some(Command::Quit),
            "s" | "scan" => Some(Command::Scan),
            "r" | "reset" => Some(Command::Reset),
            _ => {
                // 首先尝试解析为单个IP
                if let Ok(ip) = IpAddr::from_str(input) {
                    debug!("Parsed single IP, converting to /24 network");
                    match IpNetwork::new(ip, 24) {
                        Ok(network) => {
                            debug!("Converted to network: {}", network);
                            Some(Command::SetNetwork(network))
                        }
                        Err(e) => {
                            error!("Failed to convert IP to network: {}", e);
                            None
                        }
                    }
                } 
                // 然后尝试解析为网段
                else if let Ok(network) = IpNetwork::from_str(input) {
                    if network.prefix() == 32 {
                        debug!("Got /32 network, converting to /24");
                        // 如果是 /32 网络，转换为 /24
                        if let Ok(new_network) = IpNetwork::new(network.ip(), 24) {
                            debug!("Converted /32 to /24 network: {}", new_network);
                            Some(Command::SetNetwork(new_network))
                        } else {
                            error!("Failed to convert /32 to /24 network");
                            None
                        }
                    } else {
                        debug!("Successfully parsed network: {}", network);
                        Some(Command::SetNetwork(network))
                    }
                } else {
                    debug!("Invalid IP or network format");
                    None
                }
            }
        }
    }

    pub fn scroll_message_up(&mut self) {
        if self.message_scroll > 0 {
            self.message_scroll -= 1;
        }
    }

    pub fn scroll_message_down(&mut self) {
        self.message_scroll += 1;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_command() {
        let app = App::new(IpNetwork::new(IpAddr::V4([127, 0, 0, 1].into()), 24).unwrap());
        
        assert!(matches!(app.parse_command(), None));
        
        let mut app = App::new(IpNetwork::new(IpAddr::V4([127, 0, 0, 1].into()), 24).unwrap());
        app.input = "help".to_string();
        assert!(matches!(app.parse_command(), Some(Command::Help)));
        
        app.input = "192.168.1.0/24".to_string();
        if let Some(Command::SetNetwork(network)) = app.parse_command() {
            assert_eq!(network.to_string(), "192.168.1.0/24");
        } else {
            panic!("Expected SetNetwork command");
        }
    }

    #[test]
    fn test_generate_ip_list() {
        let app = App::new(IpNetwork::new(IpAddr::V4([192, 168, 1, 0].into()), 24).unwrap());
        let ips = app.generate_ip_list(&app.network);
        
        assert!(!ips.is_empty());
        assert!(ips.len() <= 254); // 不包括网络地址和广播地址
        
        // 检查第一个和最后一个IP
        assert_eq!(ips[0].address, "192.168.1.1");
        assert_eq!(ips[ips.len()-1].address, "192.168.1.254");
    }

    #[test]
    fn test_parse_command_with_single_ip() {
        let mut app = App::new(IpNetwork::new(IpAddr::V4([127, 0, 0, 1].into()), 24).unwrap());
        
        // 测试单个IP地址
        app.input = "192.168.1.1".to_string();
        if let Some(Command::SetNetwork(network)) = app.parse_command() {
            assert_eq!(network.to_string(), "192.168.1.0/24");
            assert_eq!(network.prefix(), 24);
        } else {
            panic!("Failed to parse single IP address");
        }

        // 测试 /32 网络
        app.input = "192.168.1.1/32".to_string();
        if let Some(Command::SetNetwork(network)) = app.parse_command() {
            assert_eq!(network.to_string(), "192.168.1.0/24");
            assert_eq!(network.prefix(), 24);
        } else {
            panic!("Failed to convert /32 network to /24");
        }
    }

    #[test]
    fn test_parse_command_with_invalid_input() {
        let mut app = App::new(IpNetwork::new(IpAddr::V4([127, 0, 0, 1].into()), 24).unwrap());
        
        // 测试无效输入
        app.input = "invalid_ip".to_string();
        assert!(app.parse_command().is_none());
        
        // 测试无效IP格式
        app.input = "256.256.256.256".to_string();
        assert!(app.parse_command().is_none());
    }
} 